package com.guchenbo.lssp.pool;

import com.guchenbo.lssp.core.Session;
import com.guchenbo.lssp.domain.LsspSide;
import com.guchenbo.lssp.domain.SessionKind;
import com.guchenbo.lssp.exception.GetSessionTimeoutException;
import com.guchenbo.lssp.exception.SessionException;
import com.guchenbo.lssp.support.LivyProperties;
import com.google.common.base.Preconditions;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.client.RestTemplate;

import java.io.Closeable;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.ReentrantLock;

public class SessionSource extends AbstractSessionSource implements Closeable {
    protected static final Logger log = LoggerFactory.getLogger(SessionSource.class);

    /**
     * livy session kind
     */
    private SessionKind sessionKind = SessionKind.SQL;
    /**
     * livy地址，不能以 / 结尾
     */
    private String livyHost;
    /**
     * restTemplate
     */
    private RestTemplate restTemplate;

    /**
     * 资源池当前数量
     */
    private int poolingCount = 0;
    /**
     * 资源池借出数量
     */
    private int activeCount = 0;
    private volatile SessionHolder[] sessions;
    private final ReentrantLock lock;
    private final Condition notEmpty;
    private final Condition empty;
    private final CountDownLatch initLatch;
    private CreateSessionThread createSessionThread;
    private KeepSessionThread keepSessionThread;
    private LivyHelper livyHelper;
    private LivyProperties livyProperties;
    private LsspSide lsspSide = LsspSide.CLIENT;

    public SessionSource() {
        lock = new ReentrantLock();
        notEmpty = lock.newCondition();
        empty = lock.newCondition();
        initLatch = new CountDownLatch(2);
    }

    public void init() throws SessionException {
        // 判断是否已经初始化，如果已经初始化直接跳出；如果没有初始化则继续初始化
        if (inited) {
            return;
        }
        // 加锁
        lock.lock();
        try {
            log.info("SessionSource init ...");
            // 双重校验
            if (inited) {
                return;
            }
            // 检查参数
            checkParam();
            // livy相关
            initLivy();
            // 初始化仓库数组
            sessions = new SessionHolder[maxActive];
            // 初始化几个资源
            initSession();
            // CreateThread
            createAndStartCreateThread();
            // KeepThread
            createAndStartKeepThread();
            // 保证线程先起来
            initLatch.await();
        } catch (InterruptedException e) {
            // Restore interrupted state...
            Thread.currentThread().interrupt();
            // 创建失败
            throw new SessionException("interrupt", e);
        } catch (SessionException e) {
            throw e;
        } catch (Exception e) {
            log.error("SessionSource init error", e);
            // 创建失败
            throw new SessionException(e);
        } finally {
            log.info("SessionSource[{}] inited", getName());
            //6、标记初始化完成并释放锁
            lock.unlock();
            inited = true;
        }
    }

    public PooledSession getSession() throws SessionException {
        init();
        return getSession(maxWaitMillis);
    }

    public void recycle(PooledSession pooledSession) {
        SessionHolder holder = pooledSession.getHolder();
        if (holder.isDiscard()) {
            //  如果资源坏了（使用当中坏的），这里关闭资源
            killSession(holder.getSession());
        }
        //  加锁
        lock.lock();
        try {
            //  资源没坏，就放回池子
            if (holder.isActive()) {
                putLast(holder);
            }
            //  资源坏不坏，使用数量都减少一
            holder.setActive(false);
            activeCount--;
        } finally {
            lock.unlock();
        }
    }

    @Override
    public void close() {
        log.info("SessionSource[{}] closing ...", getName());
        lock.lock();
        try {
            //  判断是否初始化了
            if (!inited) {
                return;
            }
            //  判断是否已经关闭了
            if (closed) {
                return;
            }
            closed = true;
            //  关闭池子中的资源
            for (int i = 0; i < poolingCount; i++) {
                SessionHolder holder = sessions[i];
                if (holder != null && !holder.isShared() && LsspSide.SERVER.equals(lsspSide)) {
                    // 服务端关闭，删除创建的session
                    killSession(holder.getSession());
                }
            }
            poolingCount = 0;
            activeCount = 0;
            //  关闭Create、Keep线程
            if (createSessionThread != null) {
                createSessionThread.interrupt();
            }
            if (keepSessionThread != null) {
                keepSessionThread.interrupt();
            }
            //  唤醒get线程
            notEmpty.signalAll();
        } finally {
            log.info("SessionSource[{}] closed", getName());
            lock.unlock();
        }
    }

    public int getPoolingCount() {
        lock.lock();
        try {
            return poolingCount;
        } finally {
            lock.unlock();
        }
    }

    public int getActiveCount() {
        lock.lock();
        try {
            return activeCount;
        } finally {
            lock.unlock();
        }
    }

    protected PooledSession getSession(long waitMills) throws SessionException {
        PooledSession pooledSession;
        try {
            pooledSession = getSessionInternal(waitMills);
        } catch (GetSessionTimeoutException e) {
            log.warn("get session timeout");
            throw e;
        }
        //  到这里表示池子有资源，get返回，也不用管资源是否正常，
        //  这里不检测，有地方会检测，而且Livy Session创建需要http轮训，
        //  放在这里检测太费时间，整个阻塞同步队列的都会耗时在这里，所以让调用get的这个线程自己去检测，只耗时他自己
        return pooledSession;
    }

    protected PooledSession getSessionInternal(long waitMills) throws SessionException {
        SessionHolder holder = null;
        String errorMessage = null;
        //  循环开始
        for (; ; ) {
            //  加锁
            lock.lock();
            try {
                long nanos = TimeUnit.MILLISECONDS.toNanos(waitMills);
                holder = pollLast(nanos);

                if (holder != null) {
                    //  这里可以快速判断下资源好坏
                    if (holder.isDiscard()) {
                        continue;
                    }

                    activeCount++;
                    holder.setActive(true);
                } else {
                    errorMessage = "wait second "
                                    + waitMills / 1000
                                    + ", active "
                                    + activeCount
                                    + ", maxActive "
                                    + maxActive;
                }
            } catch (InterruptedException e) {
                // Restore interrupted state...
                Thread.currentThread().interrupt();
                throw new SessionException("interrupt", e);
            } finally {
                lock.unlock();
            }
            break;
        }
        if (holder == null) {
            throw new GetSessionTimeoutException(errorMessage);
        }
        return new PooledSession(holder, livyHelper);
    }

    private void checkParam() {
        if (minIdle <= 0) {
            minIdle = initialSize;
        }
        if (StringUtils.isBlank(name)) {
            String defaultName = "SessionSource-" + sessionKind + "-" + System.identityHashCode(this);
            setName(defaultName);
        }
        Preconditions.checkArgument(getInitialSize() >= 0, "initialSize must >= 0");
        Preconditions.checkArgument(getMinIdle() > 0, "minIdle must > 0");
        Preconditions.checkArgument(getMaxActive() >= getInitialSize(), "maxActive must >= initialSize");
        Preconditions.checkArgument(getMaxActive() >= getMinIdle(), "maxActive must >= minIdle");
        Preconditions.checkArgument(getMaxWaitMillis() > 0, "maxWaitMillis must > 0");
        Preconditions.checkArgument(getTimeBetweenEvictionRunsMillis() > 0, "timeBetweenEvictionRunsMillis > 0");

        // check livy
        Preconditions.checkArgument(StringUtils.isNotBlank(livyHost), "livyHost is null");
        Preconditions.checkArgument(sessionKind != null, "kind is null");
        Preconditions.checkArgument(restTemplate != null, "restTemplate is null");
    }

    private void initLivy() {
        if (livyHelper == null) {
            livyHelper = new LivyHelper(livyHost, restTemplate, livyProperties);
        }
    }

    private void initSession() throws SessionException {
        if (initialSize == 0) {
            return;
        }
        if (shared) {
            getShareSessions();
        }
        initNewSession();
    }

    private void getShareSessions() {
        log.warn("SessionSource is in share mode, will share with same group");
        // 获取livy上面可用的session，根据lssp的kind分组，同一个分组下的session共享
        List<Session> list = livyHelper.getActiveSession(getSessionKind());
        for (Session session : list) {
            if (poolingCount < initialSize) {
                if (session != null) {
                    log.info("use share session {}", session.getId());
                    SessionHolder holder = new SessionHolder(this, session);
                    holder.setShared(true);
                    sessions[poolingCount++] = holder;
                }
            } else {
                break;
            }
        }
    }

    private void initNewSession() throws SessionException {
        boolean initSessionError = false;
        while (poolingCount < initialSize) {
            Session session = null;
            try {
                session = createRealSession();
                log.info("init new session {}", session.getId());
            } catch (SessionException e) {
                initSessionError = true;
                log.error("init sessionSource error", e);
            }
            SessionHolder holder = new SessionHolder(this, session);
            sessions[poolingCount++] = holder;
        }
        if (initSessionError && poolingCount == 0) {
            throw new SessionException("SessionSource init error");
        }
    }

    private Session createRealSession() throws SessionException {
        return livyHelper.createSession(sessionKind.getName());
    }

    private void createAndStartCreateThread() {
        String name = "LSSP-Creator-" + System.identityHashCode(this);
        createSessionThread = new CreateSessionThread(name);
        createSessionThread.start();
    }

    private void createAndStartKeepThread() {
        String name = "LSSP-Keeper-" + System.identityHashCode(this);
        keepSessionThread = new KeepSessionThread(name);
        keepSessionThread.start();
    }

    private void killSession(Session session) {
        log.warn("kill session {}", session.getId());
        livyHelper.killSession(session.getId());
    }

    private SessionHolder pollLast(long nanos) throws InterruptedException, SessionException {
        long estimate = nanos;
        //      自己被唤醒，重新判断池子，因为有可能不是被CreateThread唤醒
        //      池子还是空的，重新循环
        //  循环开始
        for (; ; ) {
            //  如果池子是空的
            if (poolingCount == 0) {
                //  唤醒CreateThread
                empty.signal();
                try {
                    //  自己阻塞 maxWait
                    estimate = notEmpty.awaitNanos(estimate);
                } catch (InterruptedException ie) {
                    // 被中断了，快速唤醒下一个等待的
                    notEmpty.signal();
                    throw ie;
                }

                //  自己被唤醒，重新判断池子，因为有可能不是被CreateThread唤醒
                //  这个判断不要去除，sonar扫描不准
                if (poolingCount == 0) {
                    if (estimate > 0) {
                        // 没有到达超时时间，继续等待
                        continue;
                    }
                    return null;
                }
            }

            decrementPoolingCount();
            SessionHolder last = sessions[poolingCount];
            sessions[poolingCount] = null;
            return last;
        }
    }

    private void put(Session session) {
        SessionHolder holder = new SessionHolder(this, session);
        lock.lock();
        try {
            putLast(holder);
        } finally {
            lock.unlock();
        }
    }

    private void putLast(SessionHolder holder) {
        sessions[poolingCount] = holder;
        incrementPoolingCount();
        notEmpty.signal();
    }

    private void decrementPoolingCount() {
        poolingCount--;
    }

    private void incrementPoolingCount() {
        poolingCount++;
    }

    public SessionKind getSessionKind() {
        return sessionKind;
    }

    public void setSessionKind(SessionKind sessionKind) {
        this.sessionKind = sessionKind;
    }

    public String getLivyHost() {
        return livyHost;
    }

    public void setLivyHost(String livyHost) {
        this.livyHost = livyHost;
    }

    public void setRestTemplate(RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
    }

    public RestTemplate getRestTemplate() {
        return restTemplate;
    }

    public void setLivyHelper(LivyHelper livyHelper) {
        this.livyHelper = livyHelper;
    }

    public LivyHelper getLivyHelper() {
        return livyHelper;
    }

    public LivyProperties getLivyProperties() {
        return livyProperties;
    }

    public void setLivyProperties(LivyProperties livyProperties) {
        this.livyProperties = livyProperties;
    }

    public class CreateSessionThread extends Thread {

        public CreateSessionThread(String name) {
            super(name);
        }

        @Override
        public void run() {
            initLatch.countDown();
            log.info("CreateSessionThread start");
            //  循环
            for (; ; ) {
                //  加锁
                lock.lock();
                try {
                    //  阻塞自己，等待唤醒
                    empty.await();
                    //  唤醒之后，再次判断是否达到资源上限，达到就再次阻塞自己
                    if (poolingCount + activeCount >= maxActive) {
                        log.warn("pool is max continue await");
                        empty.await();
                        //  再次进入循环
                        continue;
                    }
                } catch (InterruptedException e) {
                    // 响应阻断
                    log.warn("session creator interrupt");
                    // Restore interrupted state...
                    Thread.currentThread().interrupt();
                    break;
                } finally {
                    lock.unlock();
                }
                //  创建资源，加入仓库，唤醒等待者
                Session session = null;
                try {
                    session = createRealSession();
                } catch (SessionException e) {
                    //  创建失败，重新循环
                    continue;
                }
                put(session);
            }
        }

    }

    public class KeepSessionThread extends Thread {

        public KeepSessionThread(String name) {
            super(name);
        }

        @Override
        public void run() {
            initLatch.countDown();
            log.info("KeepSessionThread start");
            for (; ; ) {
                //  循环
                try {
                    //  阻塞一段时间
                    Thread.sleep(timeBetweenEvictionRunsMillis);
                    //  加锁
                    lock.lock();
                    try {
                        for (SessionHolder holder : sessions) {
                            if (holder != null && holder.getSession() != null && holder.isIdle()) {
                                // 只检查空闲的holder
                                boolean valid = livyHelper.checkSession(holder.getSession().getId());
                                if (!valid) {
                                    log.info("session {} is valid then discard", holder.getSession().getId());
                                    holder.setDiscard(true);
                                    decrementPoolingCount();
                                    sessions[poolingCount] = null;
                                }
                            }
                        }
                        //  判断是否要保持最小资源数量，需要就唤醒CreateThread
                        if (poolingCount + activeCount < minIdle) {
                            log.info("pool count {} < minIdle {} then create one", (poolingCount + activeCount),
                                            minIdle);
                            empty.signal();
                        }
                    } finally {
                        lock.unlock();
                    }
                    //  继续循环
                } catch (InterruptedException e) {
                    // 响应阻断
                    log.warn("session keeper interrupt");
                    // Restore interrupted state...
                    Thread.currentThread().interrupt();
                    break;
                }
            }

        }
    }

}
